package org.broadinstitute.hellbender.utils.iterators;

import htsjdk.samtools.util.PeekableIterator;
import org.broadinstitute.hellbender.engine.AlignmentContext;
import org.broadinstitute.hellbender.utils.SimpleInterval;
import org.broadinstitute.hellbender.utils.Utils;
import org.broadinstitute.hellbender.utils.locusiterator.IntervalAlignmentContextIterator;
import org.broadinstitute.hellbender.utils.pileup.ReadPileup;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.NoSuchElementException;

/**
 * A super-simplified/stripped-down/faster version of {@link IntervalAlignmentContextIterator} that takes a locus iterator and
 * a *single* interval, and returns an AlignmentContext for every locus in the interval. The AlignmentContexts come from
 * the provided locus iterator where coverage exists, and loci not covered by the provided locus iterator are represented
 * by an empty AlignmentContext object. AlignmentContexts generated by the provided locus iterator that lie outside of
 * the interval are discarded.
 *
 * IMPORTANT: All AlignmentContexts returned by the provided locus iterator must lie on the same contig as the
 * provided interval. It is up to the client to guarantee this (though a sanity check is done at construction time
 * on the first AlignmentContext in the client's locus iterator).
 *
 * Clients who need a more flexible and robust implementation that supports multiple intervals and contigs should use
 * {@link IntervalAlignmentContextIterator} instead.
 */
public class AllLocusIterator implements Iterator<AlignmentContext> {
    private final PeekableIterator<AlignmentContext> nestedLocusIterator;
    private final SimpleInterval interval;

    private int currentPosition;
    private AlignmentContext nextPileup;

    /**
     * @param interval The single interval over whose loci we'll be iterating
     * @param nestedLocusIterator Provider of AlignmentContexts that may lie within the interval. Must return AlignmentContexts
     *                            that are on the same contig as the provided interval.
     */
    public AllLocusIterator(final SimpleInterval interval, final Iterator<AlignmentContext> nestedLocusIterator) {
        Utils.nonNull(interval);
        Utils.nonNull(nestedLocusIterator);
        
        this.nestedLocusIterator = new PeekableIterator<>(nestedLocusIterator);
        this.interval = interval;
        this.currentPosition = interval.getStart();

        // Sanity check:
        if ( this.nestedLocusIterator.peek() != null && ! this.nestedLocusIterator.peek().getContig().equals(interval.getContig()) ) {
            throw new IllegalArgumentException("Locus iterator must be over the same contig as the interval provided");
        }

        nextPileup = advance();
    }

    @Override
    public boolean hasNext() {
        return nextPileup != null;
    }

    @Override
    public AlignmentContext next() {
        if ( ! hasNext() ) {
            throw new NoSuchElementException("next() called when there are no more items");
        }

        final AlignmentContext toReturn = nextPileup;
        nextPileup = advance();
        return toReturn;
    }

    private AlignmentContext advance() {
        // If we're out of loci, pull on the nested locus iterator until it's exhausted (caller may be relying on this),
        // and then return null
        if ( currentPosition > interval.getEnd() ) {
            while ( nestedLocusIterator.hasNext() ) {
                nestedLocusIterator.next();
            }

            return null;
        }

        // Discard AlignmentContexts from the nested locus iterator that are before our current position
        while ( nestedLocusIterator.hasNext() && nestedLocusIterator.peek().getStart() < currentPosition ) {
            nestedLocusIterator.next();
        }

        final AlignmentContext nextNestedPileup = nestedLocusIterator.peek();
        AlignmentContext toReturn;

        // No more pileups from the nested iterator? then always return an empty pileup
        if ( nextNestedPileup == null ) {
            toReturn = createEmptyAlignmentContextForPosition(currentPosition);

        // If the pileup from the nested iterator matches our current position, return it
        } else if ( nextNestedPileup.getStart() == currentPosition ) {
            toReturn = nestedLocusIterator.next();

        // Otherwise, the next pileup from our nested iterator must come after our current position,
        // so keep it around and return an empty pileup for the current position
        } else {
            toReturn = createEmptyAlignmentContextForPosition(currentPosition);
        }

        currentPosition++;
        return toReturn;
    }

    private AlignmentContext createEmptyAlignmentContextForPosition(final int position) {
        final SimpleInterval positionInterval = new SimpleInterval(interval.getContig(), position, position);
        return new AlignmentContext(positionInterval, new ReadPileup(positionInterval));
    }

    @Override
    public void remove() {
        throw new UnsupportedOperationException("remove() not supported");
    }
}
